import matplotlib.pyplot as plt
import mplfinance as mpf
from copy import copy

#解决mplfinance中文乱码问题：https://blog.csdn.net/Wilburzzz/article/details/109667926
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# 解决mplfinance绘制输出中文乱码
style = mpf.make_mpf_style(base_mpf_style='yahoo', rc={'font.family': 'SimHei'})

'''绘制K线'''
def KLine_plot(data, title, savePath,add_plot=[]):
    '''
    @param data K线数据(data可以是日K线，周K线或者月K线的数据)； type=ndarray，要求index为DatetimeIndex
    @param title 图表标题； type=str
    @param savePath 图表保存路径； type=str
    @param add_plot mpf子图; type=list
    @return 返回绘制状态
    '''

    # start_date = data.index[0]
    # end_date = data.index[-1]
    # 使用蜡烛图绘制日K线
    try:
        mpf.plot(data,
                 type='candle',
                 ylabel="price",
                 addplot=add_plot,
                 style=style,
                 title=title,
                 volume=True,
                 figratio=(5, 3),
                 mav=(5, 10, 30),
                 ylabel_lower="Volume",
                 savefig=savePath)
        return True
    except:
        print(f'关于"{title}"的K线图绘制失败')
        return False

'''绘制MA指标'''
def MA_plot(stock_data, MA_data, title, savePath):
    '''
    @param stock_data 股票数据, 其中stock_data.index是DatetimeIndex; type=dataframe
    @param MA_data 股票的MA数据; type=ndarray
    @param title 图表标题; type=str
    @param savePath 图表要保存的路径; type=str
    '''
    x = stock_data.index
    y = MA_data
    plt.plot(x,y)
    plt.title(title)
    plt.savefig(savePath)

'''绘制MACD指标'''
#mpfinance通过子图绘制macd 参考 https://blog.csdn.net/wuwei_201/article/details/108018343
def MACD_plot(macd,macdSignal,histogram,pos=2, ylabel="MACD"):
    '''
    @param macd 差离值; type=ndarray
    @param macdSignal 讯号线; type=ndarray
    @param macdhist 柱状图; type=ndarray
    @param pos 子图位置； type=int
    @param ylabel 子图y轴标签； type=str
    @return mpf子图列表
    '''

    histogram_copy = copy(histogram)
    histogram[histogram < 0] = None
    histogram_positive = histogram
    histogram_copy[histogram >= 0] = None
    histogram_negative = histogram_copy

    add_plot = [
        mpf.make_addplot(histogram_positive, type='bar', width=0.7, panel=pos, color='b',ylabel=ylabel),
        mpf.make_addplot(histogram_negative, type='bar', width=0.7, panel=pos, color='fuchsia'),
        mpf.make_addplot(macd, panel=pos, color='fuchsia', secondary_y=True),
        mpf.make_addplot(macdSignal, panel=pos, color='b', secondary_y=True)
    ]
    return add_plot

'''绘制RSI指标'''
def RSI_plot(RSI6, RSI12, RSI24, pos=2, ylabel="RSI"):
    '''
    @param RSI6 滑动窗口为6计算的RSI； type=ndarray
    @param RSI12 滑动窗口为12计算的RSI； type=ndarray
    @param RSI24 滑动窗口为24计算的RSI； type=ndarray
    @param pos 子图位置； type=int
    @param ylabel 子图y轴标签； type=str
    @return mpf子图列表
    '''
    add_plot = [
        mpf.make_addplot(RSI6, panel=pos, color='fuchsia', secondary_y=True,ylabel=ylabel),
        mpf.make_addplot(RSI12, panel=pos, color='b', secondary_y=True),
        mpf.make_addplot(RSI24, panel=pos, color='r', secondary_y=True)
    ]
    return add_plot

'''绘制Bollinger指标'''
def Bollinger_plot(H_line, M_line, L_line, pos=2, ylabel="Bollinger"):
    '''
    @param H_line 中轨线； type=ndarray
    @param M_line 上轨线； type=ndarray
    @param L_line 下轨线； type=ndarray
    @param pos 子图位置； type=int
    @param ylabel 子图y轴标签； type=str
    @return mpf子图列表
    '''
    add_plot = [
        mpf.make_addplot(H_line, panel=pos, color='fuchsia', secondary_y=True,ylabel=ylabel),
        mpf.make_addplot(M_line, panel=pos, color='b', secondary_y=True),
        mpf.make_addplot(L_line, panel=pos, color='r', secondary_y=True)
    ]
    return add_plot

'''绘制OBV指标'''
def OBV_plot(obv,pos=2,ylabel="OBV"):
    '''
    @param obv 能量潮； type=ndarray
    @param pos 子图位置； type=int
    @param ylabel 子图y轴标签； type=str
    @return mpf子图列表
    '''
    add_plot = [
        mpf.make_addplot(obv, panel=pos, color='fuchsia', secondary_y=True,ylabel=ylabel),
    ]
    return add_plot

